#!/usr/bin/env python3
"""
MCTS推理树JSONL格式可视化工具
基于vis_tree.py的可视化方法，专门处理项目生成的JSONL输出文件
"""

import json
import argparse
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from pathlib import Path

@dataclass
class TreeNode:
    """树节点数据结构"""
    node_id: str
    node_type: str
    q_value: float
    n_visits: int
    v_value: Optional[float]
    content: str
    depth: int
    action: str
    action_input: str
    final_answer: str
    children: Optional[List['TreeNode']] = None
    
    def __post_init__(self):
        if self.children is None:
            self.children = []

def parse_jsonl_file(file_path: str, max_records: int = 1) -> List[Dict[str, Any]]:
    """
    解析JSONL文件并返回指定数量的记录
    
    Args:
        file_path: JSONL文件路径
        max_records: 最大处理记录数，默认为1
        
    Returns:
        解析后的JSON记录列表
    """
    records = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            if i >= max_records:
                break
            if line.strip():
                try:
                    record = json.loads(line.strip())
                    records.append(record)
                except json.JSONDecodeError as e:
                    print(f"警告: 第{i+1}行JSON解析失败: {e}")
                    continue
    return records

def build_tree_from_rstar(rstar_data: Dict[str, Any]) -> Optional[TreeNode]:
    """
    从rstar数据构建树结构
    
    Args:
        rstar_data: rstar字段的数据
        
    Returns:
        构建的根节点
    """
    if not rstar_data:
        return None
    
    # 创建节点字典，用于快速查找
    nodes = {}
    
    # 首先创建所有节点
    for node_id, node_data in rstar_data.items():
        # 根据action字段确定节点类型
        action = node_data.get('action', '')
        if action:
            node_type = f"ACTION_{action.upper()}"
        else:
            # 如果action为空，根据节点ID判断类型
            if node_id == "0":
                node_type = "ROOT"
            else:
                node_type = "THINKING"
        
        # 计算节点深度（根据ID中的点号数量）
        depth = node_id.count('.')
        
        node = TreeNode(
            node_id=node_id,
            node_type=node_type,
            q_value=node_data.get('q_value', 0.0),
            n_visits=node_data.get('visit_count', 0),
            v_value=node_data.get('value'),
            content=node_data.get('text', ''),
            depth=depth,
            action=action,
            action_input=node_data.get('action_input', ''),
            final_answer=node_data.get('final_answer', '')
        )
        nodes[node_id] = node
    
    # 建立父子关系
    for node_id, node in nodes.items():
        if node_id == "0":
            continue  # 根节点没有父节点
        
        # 找到父节点ID
        parent_id = get_parent_id(node_id)
        if parent_id in nodes:
            parent_node = nodes[parent_id]
            if parent_node.children is None:
                parent_node.children = []
            parent_node.children.append(node)
    
    # 对每个节点的子节点按ID排序
    for node in nodes.values():
        if node.children:
            node.children.sort(key=lambda x: x.node_id)
    
    return nodes.get("0")

def get_parent_id(node_id: str) -> str:
    """
    获取节点的父节点ID
    
    Args:
        node_id: 当前节点ID，如"0.1.2"
        
    Returns:
        父节点ID，如"0.1"
    """
    parts = node_id.split('.')
    if len(parts) <= 1:
        return ""
    return '.'.join(parts[:-1])

def export_to_web(node: TreeNode, output_file: str, input_file: Optional[str] = None, record_index: int = 0):
    """导出为交互式网页显示"""
    def node_to_dict(node: TreeNode) -> Dict[str, Any]:
        """将节点转换为字典格式"""
        return {
            'id': node.node_id,
            'type': node.node_type,
            'q_value': node.q_value,
            'n_visits': node.n_visits,
            'v_value': node.v_value,
            'content': node.content,
            'depth': node.depth,
            'action': node.action,
            'action_input': node.action_input,
            'final_answer': node.final_answer,
            'children': [node_to_dict(child) for child in (node.children or [])]
        }
    
    tree_data = node_to_dict(node)
    
    # 生成动态标题
    if input_file:
        from pathlib import Path
        input_name = Path(input_file).stem
        page_title = f"{input_name} - Record {record_index} - MCTS推理树可视化"
        header_title = f"MCTS推理树可视化 - {input_name} - Record {record_index}"
    else:
        page_title = f"MCTS推理树可视化 - Record {record_index}"
        header_title = f"MCTS推理树可视化 - Record {record_index}"
    
    html_content = f"""
<!DOCTYPE html>
<html>
<head>
    <title>{page_title}</title>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <script src="https://d3js.org/d3.v7.min.js"></script>
    <style>
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        
        body {{
            font-family: 'Arial', sans-serif;
            background: #f8f9fa;
            color: #333;
            overflow: hidden;
        }}
        
        .container {{
            width: 100vw;
            height: 100vh;
            display: flex;
            flex-direction: column;
        }}
        
        .header {{
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 15px 30px;
            text-align: center;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        }}
        
        .header h1 {{
            font-size: 1.8em;
            margin: 0;
        }}
        
        .tree-container {{
            flex: 1;
            background: white;
            position: relative;
        }}
        
        .node {{
            cursor: pointer;
        }}
        
        .node circle {{
            fill: #fff;
            stroke-width: 2px;
            r: 20;
        }}
        
        /* 根据节点类型设置颜色 */
        .node.ROOT circle {{ fill: #e3f2fd; stroke: #2196f3; stroke-width: 4px; }}
        .node.THINKING circle {{ fill: #e8f5e8; stroke: #4caf50; }}
        .node.ACTION_PYTHON circle {{ fill: #fff3e0; stroke: #ff9800; }}
        .node.ACTION_SUBMIT circle {{ fill: #fce4ec; stroke: #e91e63; }}
        .node.ACTION_SEARCH circle {{ fill: #f3e5f5; stroke: #9c27b0; }}
        .node.ACTION_CALCULATOR circle {{ fill: #e0f2f1; stroke: #009688; }}
        
        .node:hover circle {{
            stroke-width: 3px;
            filter: drop-shadow(0 2px 4px rgba(0,0,0,0.2));
        }}
        
        .node.ROOT:hover circle {{
            stroke-width: 5px;
        }}
        
        .link {{
            fill: none;
            stroke: #9ecae1;
            stroke-width: 2px;
            opacity: 0.7;
        }}
        
        .node-text {{
            font-size: 11px;
            text-anchor: middle;
            pointer-events: none;
            font-weight: bold;
            fill: #333;
        }}
        
        .tooltip {{
            position: absolute;
            background: rgba(0, 0, 0, 0.9);
            color: white;
            padding: 8px 10px;
            border-radius: 6px;
            font-size: 12px;
            max-width: 600px;
            max-height: 400px;
            overflow-y: auto;
            z-index: 1000;
            display: none;
            box-shadow: 0 4px 20px rgba(0,0,0,0.3);
            line-height: 1.2;
            white-space: pre-wrap;
        }}
        
        .tooltip-title {{
            font-weight: bold;
            font-size: 14px;
            margin-bottom: 4px;
            color: #4fc3f7;
        }}
        
        .tooltip-content {{
            margin-bottom: 1px;
            line-height: 1.2;
        }}
        
        .tooltip-content strong {{
            color: #81c784;
            margin-right: 2px;
        }}
        
        .tooltip-stats {{
            display: flex;
            gap: 12px;
            margin-bottom: 4px;
            flex-wrap: wrap;
        }}
        
        .tooltip-stat {{
            font-size: 11px;
            line-height: 1.1;
        }}
        
        .tooltip-content-text {{
            margin-top: 4px;
            padding-left: 4px;
            font-size: 11px;
            line-height: 1.2;
            border-left: 2px solid #4fc3f7;
            padding-left: 6px;
        }}
        
        .legend {{
            position: absolute;
            top: 20px;
            right: 20px;
            background: rgba(255,255,255,0.95);
            padding: 15px;
            border-radius: 8px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            font-size: 12px;
            max-width: 250px;
        }}
        
        .legend-title {{
            font-weight: bold;
            margin-bottom: 10px;
            color: #333;
        }}
        
        .legend-item {{
            display: flex;
            align-items: center;
            margin-bottom: 8px;
        }}
        
        .legend-color {{
            width: 16px;
            height: 16px;
            border-radius: 50%;
            margin-right: 8px;
            border: 2px solid;
        }}
        
        .legend-color.ROOT {{ background: #e3f2fd; border-color: #2196f3; }}
        .legend-color.THINKING {{ background: #e8f5e8; border-color: #4caf50; }}
        .legend-color.ACTION_PYTHON {{ background: #fff3e0; border-color: #ff9800; }}
        .legend-color.ACTION_SUBMIT {{ background: #fce4ec; border-color: #e91e63; }}
        .legend-color.ACTION_SEARCH {{ background: #f3e5f5; border-color: #9c27b0; }}
        .legend-color.ACTION_CALCULATOR {{ background: #e0f2f1; border-color: #009688; }}
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>{header_title}</h1>
        </div>
        
        <div class="tree-container">
            <div class="legend">
                <div class="legend-title">节点类型</div>
                <div class="legend-item">
                    <div class="legend-color ROOT"></div>
                    <span>根节点 (ROOT)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color THINKING"></div>
                    <span>思考节点 (THINKING)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color ACTION_PYTHON"></div>
                    <span>Python执行 (ACTION_PYTHON)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color ACTION_SUBMIT"></div>
                    <span>提交答案 (ACTION_SUBMIT)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color ACTION_SEARCH"></div>
                    <span>搜索 (ACTION_SEARCH)</span>
                </div>
                <div class="legend-item">
                    <div class="legend-color ACTION_CALCULATOR"></div>
                    <span>计算器 (ACTION_CALCULATOR)</span>
                </div>
            </div>
            
            <div class="tooltip" id="tooltip"></div>
            <svg id="tree-svg"></svg>
        </div>
    </div>

    <script>
        const treeData = {json.dumps(tree_data, ensure_ascii=False, indent=2)};
        
        // 设置SVG尺寸
        const container = document.querySelector('.tree-container');
        const width = container.clientWidth;
        const height = container.clientHeight;
        
        const svg = d3.select('#tree-svg')
            .attr('width', width)
            .attr('height', height);
        
        const g = svg.append('g');
        
        // 创建树布局
        const tree = d3.tree()
            .size([width - 100, height - 100])
            .separation((a, b) => {{
                return a.parent === b.parent ? 1.5 : 2;
            }});
        
        // 转换数据
        const root = d3.hierarchy(treeData);
        tree(root);
        
        // 调整位置
        root.each(d => {{
            d.x += 50;
            d.y += 50;
        }});
        
        // 绘制连线
        g.selectAll('.link')
            .data(root.links())
            .enter()
            .append('path')
            .attr('class', 'link')
            .attr('d', d3.linkVertical()
                .x(d => d.x)
                .y(d => d.y)
            );
        
        // 绘制节点
        const nodes = g.selectAll('.node')
            .data(root.descendants())
            .enter()
            .append('g')
            .attr('class', d => `node ${{d.data.type}}`)
            .attr('transform', d => `translate(${{d.x}},${{d.y}})`);
        
        nodes.append('circle');
        
        nodes.append('text')
            .attr('class', 'node-text')
            .attr('dy', 4)
            .text(d => d.data.id);
        
        // 添加工具提示
        const tooltip = d3.select('#tooltip');
        let tooltipTimeout;
        
        // 鼠标悬停显示完整信息
        nodes.on('mouseover', function(event, d) {{
            clearTimeout(tooltipTimeout);
            
            // 构建工具提示内容
            let content = `
                <div class="tooltip-title">${{d.data.id}} (${{d.data.type}})</div>
                <div class="tooltip-stats">
                    <div class="tooltip-stat"><strong>Q值:</strong>${{d.data.q_value.toFixed(2)}}</div>
                    <div class="tooltip-stat"><strong>访问次数:</strong>${{d.data.n_visits}}</div>
                    <div class="tooltip-stat"><strong>V值:</strong>${{d.data.v_value !== null ? d.data.v_value.toFixed(2) : 'None'}}</div>
                </div>
            `;
            
            if (d.data.action) {{
                content += `<div class="tooltip-content"><strong>动作:</strong>${{d.data.action}}</div>`;
            }}
            
            if (d.data.action_input) {{
                content += `<div class="tooltip-content"><strong>动作输入:</strong>${{d.data.action_input}}</div>`;
            }}
            
            if (d.data.final_answer) {{
                content += `<div class="tooltip-content"><strong>最终答案:</strong>${{d.data.final_answer}}</div>`;
            }}
            
            if (d.data.content) {{
                content += `<div class="tooltip-content-text">${{d.data.content}}</div>`;
            }}
            
            tooltip.html(content)
                .style('display', 'block')
                .style('left', Math.min(event.pageX + 10, window.innerWidth - 620) + 'px')
                .style('top', Math.min(event.pageY - 10, window.innerHeight - 420) + 'px');
        }})
        .on('mouseout', function() {{
            tooltipTimeout = setTimeout(() => {{
                if (tooltip.style('display') === 'block') {{
                    const tooltipElement = document.getElementById('tooltip');
                    if (!tooltipElement.matches(':hover')) {{
                        tooltip.style('display', 'none');
                    }}
                }}
            }}, 300);
        }});
        
        // 鼠标进入tooltip时保持显示
        tooltip.on('mouseenter', function() {{
            clearTimeout(tooltipTimeout);
        }})
        .on('mouseleave', function() {{
            tooltip.style('display', 'none');
        }});
        
        // 添加缩放和拖拽
        const zoom = d3.zoom()
            .scaleExtent([0.1, 3])
            .on('zoom', function(event) {{
                g.attr('transform', event.transform);
            }});
        
        svg.call(zoom);
        
        // 适应窗口大小变化
        window.addEventListener('resize', function() {{
            const newWidth = container.clientWidth;
            const newHeight = container.clientHeight;
            svg.attr('width', newWidth).attr('height', newHeight);
            
            tree.size([newWidth - 100, newHeight - 100]);
            tree(root);
            
            root.each(d => {{
                d.x += 50;
                d.y += 50;
            }});
            
            g.selectAll('.link')
                .attr('d', d3.linkVertical()
                    .x(d => d.x)
                    .y(d => d.y)
                );
            
            g.selectAll('.node')
                .attr('transform', d => `translate(${{d.x}},${{d.y}})`);
        }});
    </script>
</body>
</html>
    """
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(html_content)
    
    print(f"网页可视化已保存到: {output_file}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='MCTS推理树JSONL格式可视化工具')
    parser.add_argument('input_file', help='输入的JSONL文件路径')
    parser.add_argument('-o', '--output', default=None, 
                       help='输出HTML文件路径 (默认: 与输入文件同名但扩展名为.html)')
    parser.add_argument('-n', '--num-records', type=int, default=1,
                       help='处理的记录数量 (默认: 1)')
    
    args = parser.parse_args()
    
    # 检查输入文件
    input_path = Path(args.input_file)
    if not input_path.exists():
        print(f"错误: 输入文件 {args.input_file} 不存在")
        return
    
    print(f"正在解析文件: {args.input_file}")
    print(f"处理记录数量: {args.num_records}")
    
    # 解析JSONL文件
    try:
        records = parse_jsonl_file(args.input_file, args.num_records)
        if not records:
            print("错误: 没有找到有效的JSON记录")
            return
        
        print(f"成功解析 {len(records)} 条记录")
        
        # 处理每条记录
        for i, record in enumerate(records):
            if 'rstar' not in record:
                print(f"警告: 第{i+1}条记录没有rstar字段，跳过")
                continue
            
            # 构建树结构
            root = build_tree_from_rstar(record['rstar'])
            if not root:
                print(f"警告: 第{i+1}条记录无法构建树结构，跳过")
                continue
            
            print(f"第{i+1}条记录 - 根节点: {root.node_id}, 总节点数: {len(record['rstar'])}")
            
            # 确定输出文件名
            if args.output is None:
                if len(records) == 1:
                    output_file = str(input_path.with_suffix('.html'))
                else:
                    output_file = str(input_path.with_suffix(f'_record_{i}.html'))
            else:
                if len(records) == 1:
                    output_file = args.output
                else:
                    output_path = Path(args.output)
                    output_file = str(output_path.with_stem(f"{output_path.stem}_record_{i}"))
            
            # 生成可视化
            export_to_web(root, output_file, args.input_file, i)
            
    except Exception as e:
        print(f"错误: 处理文件时发生异常: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 